package hn.cch.ai.alibaba.graph.configuration;

import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncNodeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.node.LlmNode;
import com.alibaba.cloud.ai.graph.observation.GraphObservationLifecycleListener;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import hn.cch.ai.alibaba.graph.node.ExpanderNode;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;

@Configuration
public class GraphConfiguration {

    private static final Logger logger = LoggerFactory.getLogger(GraphConfiguration.class);


    @Bean
    public StateGraph stateGraph(ChatClient.Builder chatClientBuilder) throws GraphStateException {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();

            // 用户输入
            keyStrategyHashMap.put("query", new ReplaceStrategy());
            keyStrategyHashMap.put("expander_number", new ReplaceStrategy());
            keyStrategyHashMap.put("expander_content", new ReplaceStrategy());
            return keyStrategyHashMap;
        };

        StateGraph stateGraph = new StateGraph(keyStrategyFactory)
            .addNode("expander", AsyncNodeAction.node_async(new ExpanderNode(chatClientBuilder)))
            .addEdge(StateGraph.START, "expander").addEdge("expander", StateGraph.END);

        // 添加 PlantUML 打印
        GraphRepresentation representation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML, "expander flow");
        logger.info("\n=== expander UML Flow ===");
        logger.info("\n{}", representation.content());
        logger.info("==================================\n");
        return stateGraph;
    }

    // @Bean
    // public StateGraph workflowGraph(ChatModel chatModel) throws GraphStateException {
    //     // 1. 创建ChatClient
    //     // ChatClient chatClient = ChatClient.builder(chatModel).build();
    //     ChatClient chatClient = ChatClient.builder(chatModel)
    //         .defaultAdvisors(new SimpleLoggerAdvisor())
    //         .build();
    //
    //
    //     // 2. 定义节点
    //     LlmNode welcomeNode = LlmNode.builder()
    //         .chatClient(chatClient)
    //         .systemPromptTemplate("你是一个友好的助手")
    //         .userPromptTemplate("欢迎用户：{input}")
    //         .outputKey("welcome_message")
    //         .build();
    //
    //     // 3. 定义状态策略
    //     KeyStrategyFactory stateFactory = () -> {
    //         Map<String, KeyStrategy> strategies = new HashMap<>();
    //         strategies.put("input", new ReplaceStrategy());
    //         strategies.put("welcome_message", new ReplaceStrategy());
    //         return strategies;
    //     };
    //
    //     // 4. 构建工作流
    //     return new StateGraph("我的第一个工作流", stateFactory)
    //         .addNode("welcome", AsyncNodeAction.node_async(welcomeNode))
    //         .addEdge(StateGraph.START, "welcome")
    //         .addEdge("welcome", StateGraph.END);
    // }
    //
    // @Bean
    // public CompiledGraph compiledGraph(StateGraph stateGraph,
    //                                    ObservationRegistry observationRegistry) throws GraphStateException {
    //     return stateGraph.compile(CompileConfig.builder()
    //         .withLifecycleListener(new GraphObservationLifecycleListener(observationRegistry))
    //         .build());
    // }
}
